package com.ics.atable.chat.config;

import com.alibaba.cloud.ai.graph.KeyStrategy;
import com.alibaba.cloud.ai.graph.KeyStrategyFactory;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.graph.state.strategy.ReplaceStrategy;
import com.ics.atable.chat.node.SqlMetricNormalizeNode;
import com.ics.atable.chat.service.MetricAliasService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import static com.alibaba.cloud.ai.graph.StateGraph.END;
import static com.alibaba.cloud.ai.graph.StateGraph.START;

import java.util.HashMap;

import static com.alibaba.cloud.ai.graph.action.AsyncNodeAction.node_async;

/**
 * nl2sql工作流
 *
 * @author haohongbin
 */
@Slf4j
@Configuration
public class Nl2sqlConfiguration {

    @Autowired
    private ChatModel chatModel;

    /**
     * SQL 统计项替换工作流名称
     * <p>
     * 该工作流用于将输入 SQL 中的统计项/口径等别名统一替换为标准名称。
     * Graph 拓扑：START -> SQL_REPLACE_NODE -> END
     */
    public static final String SQL_REPLACE_GRAPH = "SQL_REPLACE_GRAPH";
    public static final String SQL_REPLACE_NODE = "SQL_REPLACE_NODE";

    /**
     * 状态键名称约定：
     * replace - 替换规则集合（支持 Map 或 List[{from,to}]）
     * output - 输出的替换后 SQL 字符串
     */
    public static final String REPLACE = "replace";
    public static final String OUTPUT = "output";

    /**
     * 创建 ChatClient Bean。
     *
     * @return ChatClient 实例
     */
    @Bean
    @Qualifier("sqlMetricReplaceGraph")
    public ChatClient sqlMetricReplaceGraphChatClient() {
        return ChatClient.builder(chatModel).build();
    }

    /**
     * 创建 SQL 统计项替换工作流的 StateGraph Bean。
     *
     * @param chatClient         ChatClient（当前节点未直接使用，预留扩展：例如未命中规则时使用 LLM 兜底）
     * @param metricAliasService 规则替换服务，负责将 replace 规则应用到 input SQL 上
     * @return StateGraph 已装配的工作流图
     * @throws GraphStateException 构建图过程中的异常
     */
    @Bean
    public StateGraph sqlMetricReplaceGraph(@Qualifier("sqlMetricReplaceGraph") ChatClient chatClient, MetricAliasService metricAliasService) throws GraphStateException {
        KeyStrategyFactory keyStrategyFactory = () -> {
            HashMap<String, KeyStrategy> keyStrategyHashMap = new HashMap<>();
            // 声明需要在状态中读写的键，并使用 ReplaceStrategy 进行覆盖写入
            keyStrategyHashMap.put(REPLACE, new ReplaceStrategy());
            keyStrategyHashMap.put(OUTPUT, new ReplaceStrategy());
            return keyStrategyHashMap;
        };

        // 仅一个节点：执行批量替换逻辑
        StateGraph stateGraph = new StateGraph(SQL_REPLACE_GRAPH, keyStrategyFactory)
                .addNode(SQL_REPLACE_NODE, node_async(new SqlMetricNormalizeNode(metricAliasService)))
                .addEdge(START, SQL_REPLACE_NODE)
                .addEdge(SQL_REPLACE_NODE, END);

        return stateGraph;
    }
}
